package midend;

import midend.IR_Value.IR_BaseBlock;
import midend.IR_Value.IR_FuncBlock;
import midend.IR_Instruction.*;
//import backend.alloctor.GraphColoring_RegisterAllocator;
import java.util.*;

/*
* 该类用于对CFG进行控制流分析，求解支配边界，方便phi指令插入
* */
public class Dom_Analysis {
    // 待分析的函数块
    private final IR_FuncBlock functionBlock;
    // 函数块中的基本块
    private final List<IR_BaseBlock> basicBlocks;
    // 记录每个基本块的支配节点
    private final Map<IR_BaseBlock, Set<IR_BaseBlock>> dominators;
    // 记录每个基本块的直接支配节点
    private final Map<IR_BaseBlock, IR_BaseBlock> immediaDonminators;
    // 记录每个基本块的支配边界
    private final Map<IR_BaseBlock, Set<IR_BaseBlock>> dominanceFrontiers;
    // 前驱、后继
    // GraphColoring_RegisterAllocator graphColoring_registerAllocator = new GraphColoring_RegisterAllocator();

    Map<IR_BaseBlock, List<IR_BaseBlock>> predecessors = new HashMap<>();

    // 主方法
    public Dom_Analysis(IR_FuncBlock functionBlock) {
        this.functionBlock = functionBlock;
        this.basicBlocks = functionBlock.get_bb_List();
        this.dominators = new HashMap<>();
        this.immediaDonminators = new HashMap<>();
        this.dominanceFrontiers = new HashMap<>();
        this.predecessors = calcuPredecessors(functionBlock);
        //初始化各节点的支配者为全部节点
        for (IR_BaseBlock b : basicBlocks) {
            dominators.put(b, new HashSet<>(basicBlocks));
        }
        calcuDominators();
        calcuImmediaDominators();
        calcuDominanceFrontiers();
    }

    public List<IR_BaseBlock> get_Immediate_Successors(IR_BaseBlock block) {
        List<IR_BaseBlock> successors = new ArrayList<>();
        List<Instruction> instructions=block.get_Instruction_List();
        
        // 如果指令列表为空，返回空的后继列表
        if(instructions.isEmpty()) {
            return successors;
        }
        
        // 遍历所有指令，找到所有分支指令
        for (int i = instructions.size() - 1; i >= 0; i--) {
            Instruction instruction = instructions.get(i);
            if (instruction instanceof Ret_Instruction) {
                return successors;
            }
            if (instruction instanceof Branch_Instruction) {
                Branch_Instruction branch = (Branch_Instruction) instruction;
                IR_BaseBlock block1 = branch.get_block1();
                if (!successors.contains(block1)) {
                    successors.add(block1);
                }
                if (branch.get_block2() != null) {
                    IR_BaseBlock block2 = branch.get_block2();
                    if (!successors.contains(block2)) {
                        successors.add(block2);
                    }
                }
            }
        }
        
        // 更新前驱关系
        List<IR_BaseBlock> temp = new ArrayList<>();
        temp.add(block);
        for (IR_BaseBlock successor : successors) {
            if (predecessors.get(successor) == null) {
                predecessors.put(successor, new ArrayList<>());
            }
            predecessors.get(successor).addAll(temp);
        }
        
        return successors;
    }


    // 计算各基本块前驱
    public Map<IR_BaseBlock, List<IR_BaseBlock>> calcuPredecessors(IR_FuncBlock funcBlock) {
        Map<IR_BaseBlock,List<IR_BaseBlock>> predecessors = new HashMap<>();
        // 初始化
        for (IR_BaseBlock bb : funcBlock.get_bb_List()) {
            predecessors.put(bb, new ArrayList<>());
        }
        // 遍历函数里的每个基本块，利用后继列表获取前驱列表
        for (IR_BaseBlock bb : funcBlock.get_bb_List()) {
            for (IR_BaseBlock successor : this.get_Immediate_Successors(bb)) {
                predecessors.get(successor).add(bb);
            }
        }
        return predecessors;
    }

    // 计算各基本块的支配节点
    private void calcuDominators() {
        // 初始入口信息
        IR_BaseBlock entry = functionBlock.get_bb_List().get(0);
        dominators.get(entry).clear(); // 清空入口块的支配者，避免历史结果遗留
        dominators.get(entry).add(entry);

        boolean changed;
        do {
            changed = false;
            // 这里要用逆序后序遍历序列是为了满足从CFG入口进入执行的顺序
            List<IR_BaseBlock> reversePostOrder = getReversePostOrder(functionBlock);
            for (IR_BaseBlock block : reversePostOrder) {
                if (block == entry) continue;
                Set<IR_BaseBlock> updatedDominators = calcuUpdatedDominators(block);
                if(!updatedDominators.equals(dominators.get(block))) {
                    dominators.put(block, updatedDominators);
                    changed = true;
                }
            }
        } while(changed);
    }

    // 一个节点的支配者是它所有前驱节点的支配者的交集加上它本身
    private Set<IR_BaseBlock> calcuUpdatedDominators(IR_BaseBlock bb) {
        Set<IR_BaseBlock> newSet = new HashSet<>();
        List<IR_BaseBlock> preds = predecessors.get(bb);
        if (!preds.isEmpty()) {
            newSet.addAll(dominators.get(preds.get(0)));
            for (IR_BaseBlock preBlock : preds) {
                newSet = union(newSet, dominators.get(preBlock));
            }
        }
        newSet.add(bb);
        return newSet;
    }

    // 计算各基本块直接支配者
    private void calcuImmediaDominators() {
        // 遍历除入口节点外的所有基本块节点
        for (IR_BaseBlock bb : functionBlock.get_bb_List()) {
            if (bb == functionBlock.get_bb_List().get(0)) continue;
            IR_BaseBlock idom = null;
            for (IR_BaseBlock dom : dominators.get(bb)) {
                if (dom == bb) continue;
                HashSet<IR_BaseBlock> tmpDomSet = new HashSet<>(dominators.get(bb));
                tmpDomSet.remove(bb);
                //tmpDomSet是去掉自己的所有支配者
                //如果支配者A中，某支配者的所有支配者是A去掉自己的支配者集合，则最近？  // 数学理论基础：DOM(b) = {b} ∪ DOM(idom(b))
                if (idom == null && dominators.get(dom).equals(tmpDomSet)) {
                    idom = dom;
                    bb.setParentBlock(idom);
                    idom.addChildBlock(bb);
                }
            }
            immediaDonminators.put(bb,idom);

        }

    }

    //
    private Set<IR_BaseBlock> union(Set<IR_BaseBlock> set1, Set<IR_BaseBlock> set2)
    {
        Set<IR_BaseBlock> largerSet = set1.size()>=set2.size()?set1:set2;
        Set<IR_BaseBlock> lessSet = set1.size()>=set2.size()?set2:set1;
        Set<IR_BaseBlock> unionSet = new HashSet<>();
        Map<IR_BaseBlock, Boolean> map = new HashMap<>();
        for (IR_BaseBlock bb : largerSet){
            map.put(bb,true);
        }
        for (IR_BaseBlock bb : lessSet) {
            if (map.containsKey(bb)) {
                unionSet.add(bb);
            }
        }
        return unionSet;
    }

    // 后序遍历DAG
    private void dfs(IR_BaseBlock bb, List<IR_BaseBlock> postOrder, Set<IR_BaseBlock> visited) {
        boolean isvisited = visited.contains(bb);
        if (!isvisited) {
            visited.add(bb);
            for (IR_BaseBlock successor : this.get_Immediate_Successors(bb)) {
                dfs(successor, postOrder, visited);
            }
            postOrder.add(bb);
        }
    }


    // 获取反向后序序列
    // 为什么使用反向后序序列，因为它更符合从入口块开始的逻辑顺序，可以复用前驱节点的支配节点
    private List<IR_BaseBlock> getReversePostOrder(IR_FuncBlock functionBlock) {
        List<IR_BaseBlock> postOrder = getPostOrder(functionBlock);
        Collections.reverse(postOrder);
        return postOrder;
    }

    //计算各节点的支配边界
    public void calcuDominanceFrontiers() {
        // 初始化支配边界集合
        for (IR_BaseBlock bb : basicBlocks) {
            dominanceFrontiers.put(bb, new HashSet<>());
        }

        for (IR_BaseBlock bb : basicBlocks) {
            List<IR_BaseBlock> preds = predecessors.get(bb);
            for (IR_BaseBlock climber : preds) {
                while (climber!=null && !dominators.get(bb).contains(climber)) {
                    dominanceFrontiers.get(climber).add(bb);
                    climber = immediaDonminators.get(climber);
                }
            }
        }
    }

    public List<IR_BaseBlock> getPostOrder(IR_FuncBlock functionBlock) {
        Set<IR_BaseBlock> visited = new HashSet<>();
        List<IR_BaseBlock> postOrder = new ArrayList<>();
        IR_BaseBlock entry = functionBlock.get_bb_List().get(0);
        dfs(entry, postOrder,visited);
        return postOrder;
    }

    // 以下函数逻辑较为简单，主要功能是为外部提供类的访问接口
    public boolean dominates(IR_BaseBlock dominator, IR_BaseBlock dominatedBlock) {
        if (dominators.get(dominatedBlock) == null){
            return false;
        }
        return dominators.get(dominatedBlock).contains(dominator);
    }

    public Map<IR_BaseBlock, Set<IR_BaseBlock>> getDominators() { return dominators; }

    public Map<IR_BaseBlock, IR_BaseBlock> getImmediaDominators() { return immediaDonminators; }

    public Map<IR_BaseBlock, Set<IR_BaseBlock>> getDominanceFrontiers() { return dominanceFrontiers; }

}
